package com.klw.business.utils;


import com.klw.business.domain.KlwMlFeatures;
import weka.classifiers.Classifier;
import weka.core.*;

import java.io.File;
import java.util.List;

/**
 * 机器学习模型工具类
 * 负责特征数据与WEKA格式转换、模型序列化等操作
 */
public class MLModelUtils {

    /**
     * 创建WEKA所需的数据集结构
     * @return 包含特征属性的数据集结构
     */
    public static Instances createInstancesStructure() {
        // 定义特征属性
        FastVector attributes = new FastVector();

        // 添加特征属性（与数据库表字段对应）
        attributes.addElement(new Attribute("avg_score"));          // 平均成绩
        attributes.addElement(new Attribute("fail_count"));         // 挂科门数
        attributes.addElement(new Attribute("absent_score"));       // 考勤分数
        attributes.addElement(new Attribute("behavior_score"));     // 行为分数
        attributes.addElement(new Attribute("credit_score"));       // 学分总和

        // 添加类别属性（预警等级：0-3）
        FastVector classValues = new FastVector();
        classValues.addElement("0");
        classValues.addElement("1");
        classValues.addElement("2");
        classValues.addElement("3");
        attributes.addElement(new Attribute("warning_level", classValues));

        // 创建空数据集（初始容量1000条记录）
        Instances structure = new Instances("AcademicWarningData", attributes, 1000);
        // 设置最后一个属性为类别属性
        structure.setClassIndex(structure.numAttributes() - 1);

        return structure;
    }

    /**
     * 将特征实体转换为WEKA的Instance对象
     * @param feature 特征数据实体
     * @param structure 数据集结构
     * @return WEKA的Instance对象
     */
    public static Instance convertToInstance(KlwMlFeatures feature, Instances structure) {
        Instance instance = new DenseInstance(structure.numAttributes());

        // 设置特征值（顺序需与createInstancesStructure中定义的一致）
        instance.setValue(0, feature.getAvgScore() != null ? feature.getAvgScore().doubleValue() : 0.0);
        instance.setValue(1, feature.getFailCount() != null ? feature.getFailCount() : 0);
        instance.setValue(2, feature.getAbsentScore() != null ? feature.getAbsentScore() : 0);
        instance.setValue(3, feature.getBehaviorScore() != null ? feature.getBehaviorScore() : 0);
        instance.setValue(4, feature.getCreditScore() != null ? feature.getCreditScore().doubleValue() : 0.0);

        // 设置数据集引用
        instance.setDataset(structure);

        return instance;
    }

    /**
     * 将特征列表转换为WEKA的Instances数据集
     * @param features 特征列表
     * @param structure 数据集结构
     * @return WEKA的Instances数据集
     */
    public static Instances convertToInstances(List<KlwMlFeatures> features, Instances structure) {
        Instances instances = new Instances(structure);

        for (KlwMlFeatures feature : features) {
            instances.add(convertToInstance(feature, structure));
        }

        return instances;
    }

    /**
     * 保存模型到文件
     * @param classifier 分类器模型
     * @param modelPath 模型文件路径
     * @throws Exception 保存过程中的异常
     */
    public static void saveModel(Classifier classifier, String modelPath) throws Exception {
        // 创建父目录（如果不存在）
        File modelFile = new File(modelPath);
        File parentDir = modelFile.getParentFile();
        if (parentDir != null && !parentDir.exists()) {
            boolean created = parentDir.mkdirs();
            if (!created) {
                throw new Exception("无法创建模型目录: " + parentDir.getAbsolutePath());
            }
            System.out.println("模型目录已创建: " + parentDir.getAbsolutePath());
        }

        // 序列化模型
        SerializationHelper.write(modelPath, classifier);
        System.out.println("模型已保存到: " + modelPath);
    }

    /**
     * 从文件加载模型
     * @param modelPath 模型文件路径
     * @return 加载的分类器模型
     * @throws Exception 加载过程中的异常
     */
    public static Classifier loadModel(String modelPath) throws Exception {
        File modelFile = new File(modelPath);
        if (!modelFile.exists()) {
            throw new Exception("模型文件不存在: " + modelPath);
        }
        if (!modelFile.canRead()) {
            throw new Exception("无法读取模型文件: " + modelPath);
        }
        
        // 反序列化模型
        Classifier classifier = (Classifier) SerializationHelper.read(modelPath);
        System.out.println("模型已从以下路径加载: " + modelPath);
        return classifier;
    }

    /**
     * 为实例设置类别标签（用于训练数据）
     * @param instance 实例对象
     * @param warningLevel 预警等级
     */
    public static void setClassLabel(Instance instance, Integer warningLevel) {
        if (instance != null && instance.dataset() != null && warningLevel != null) {
            int classIndex = instance.dataset().classIndex();
            Attribute classAttr = instance.dataset().attribute(classIndex);
            int labelIndex = classAttr.indexOfValue(warningLevel.toString());
            if (labelIndex >= 0) {
                instance.setClassValue(labelIndex);
            }
        }
    }
}
